# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import ujson as json
import functools

from coin.base.book.types import TradeSide, create_market_trade_event
from coin.exchange.base.ws.handler import WebsocketHandler
from coin.exchange.hitbtc_v2.book_builder import HitbtcBookBuilder
from coin.exchange.hitbtc_v2.kr_rest.product import HitbtcProduct


class HandlerBase(WebsocketHandler):
  BookBuilderType = None
  TraitsType = None

  def on_ws_depth(self, record, queue_data, topic_data):
    data = record.data
    if not isinstance(data, dict):
      data = json.loads(data.tobytes())

    assert 'symbol' in data['params']
    symbol = data['params']['symbol']
    product = self.TraitsType.get_product_from_symbol(symbol)
    native_symbol = product.native_symbol
    if not self.is_subscribed('book', native_symbol):
      return

    assert 'method' in data
    method = data['method']
    if method == 'snapshotOrderbook':
      book = self.BookBuilderType(native_symbol, book_builder_observer=self._book_builder_observer)
      bids = data['params']['bid']
      asks = data['params']['ask']
      book.snapshot(bids, asks)
      self._books[native_symbol] = book
    elif method == 'updateOrderbook':
      if native_symbol not in self._books:
        return
      book = self._books[native_symbol]
      bids = data['params']['bid']
      asks = data['params']['ask']
      book.update_diff(bids, asks)
    else:
      raise ValueError(data)
    book.update_timestamp(record.timestamp)
    book.update_topic_id(record.topic_id)
    self.ready = True
    if self._feed_checker:
      self._feed_checker.on_book(product=product, book=book, timestamp=record.timestamp)
    self.publish_event('book', book.symbol, book)

  def on_ws_trade(self, record, queue_data, topic_data):
    data = record.data
    if not isinstance(data, dict):
      data = json.loads(data.tobytes())

    assert 'symbol' in data['params']
    symbol = data['params']['symbol']
    product = self.TraitsType.get_product_from_symbol(symbol)
    native_symbol = product.native_symbol
    if not self.is_subscribed('trade', native_symbol):
      return

    assert 'method' in data
    method = data['method']
    if method == 'snapshotTrades':
      # A lot of old trades at the beginning. Ignore them.
      return

    data = data['params']['data']
    for item in data:
      price = float(item['price'])
      qty = float(item['quantity'])
      if item['side'] == 'buy':
        side = TradeSide.TRADE_BUY_SIDE
      elif item['side'] == 'sell':
        side = TradeSide.TRADE_SELL_SIDE
      else:
        raise ValueError('Unknown side: %s', data)
      trade_event = create_market_trade_event(record.timestamp,
                                              native_symbol,
                                              price,
                                              qty,
                                              side,
                                              topic_id=record.topic_id)
      if self._feed_checker:
        self._feed_checker.on_trade(product=product,
                                    trade_event=trade_event,
                                    timestamp=record.timestamp)
      self.publish_event('trade', native_symbol, trade_event)


class HitbtcTraits(object):
  @classmethod
  @functools.lru_cache(maxsize=32)
  def get_product_from_symbol(cls, symbol):
    return HitbtcProduct.FromStrNativeProduct(symbol)


class HitbtcHandler(HandlerBase):
  BookBuilderType = HitbtcBookBuilder
  TraitsType = HitbtcTraits
